import os
import sqlite3
import shutil
from pathlib import Path
import sys
import argparse

def check_and_connect_db(db_path):
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row  # 设置返回字典格式的结果
    cursor = conn.cursor()
    return conn, cursor

def safe_copy(source_file, destination_dir):
    """安全复制文件，处理文件名冲突"""
    destination_file = Path(destination_dir) / source_file.name
    counter = 1
    while destination_file.exists():
        new_filename = f"{source_file.stem}({counter}){source_file.suffix}"
        destination_file = destination_file.with_name(new_filename)
        counter += 1
    shutil.copy2(str(source_file), str(destination_file))
    return destination_file.name

def validate_folders(db_directory, cursor):
    """在程序启动时验证所有需要的目标文件夹是否存在"""
    cursor.execute("SELECT DISTINCT aim_zip FROM photo_move WHERE deal_flag=1")
    aim_zips = cursor.fetchall()

    missing_folders = []
    for aim_zip_row in aim_zips:
        aim_zip = aim_zip_row['aim_zip']
        folder_path = (Path(db_directory) / aim_zip).with_suffix('')
        if not folder_path.exists():
            missing_folders.append(folder_path)

    if missing_folders:
        print(f"错误: 以下目标文件夹不存在，请检查流程是否正确执行:")
        for folder in missing_folders:
            print(folder)
        sys.exit(1)
    else:
        print("所有目标文件夹都已存在，准备继续处理文件...")

def process_record(record, db_directory, scan_directory, db_conn, cursor):
    aim_zip = record['aim_zip']
    print(f"处理文件: {record['file_name']} 到 {aim_zip}")
    
    # 根据aim_zip在db_directory下找到或创建相应的文件夹
    folder_path = (Path(db_directory) / aim_zip).with_suffix('')

    # 复制文件到目标文件夹
    source_file = Path(scan_directory) / record['file_path'] / record['file_name']
    if not source_file.exists():
        print(f"错误: 源文件 {source_file} 不存在")
        return

    copied_file_name = safe_copy(source_file, folder_path)
    print(f"已复制文件: {copied_file_name}")

    # 更新数据库记录
    cursor.execute("UPDATE photo_move SET deal_flag=2 WHERE id=?", (record['id'],))
    db_conn.commit()

def main(db_directory, scan_directory):
    db_path = str(Path(db_directory) / 'file_hashes.db')
    conn, cursor = check_and_connect_db(db_path)

    # 在程序启动时验证所有需要的目标文件夹是否存在
    validate_folders(db_directory, cursor)

    # 查询photo_move表中deal_flag=1的记录
    cursor.execute("SELECT * FROM photo_move WHERE deal_flag=1")
    rows = cursor.fetchall()

    if not rows:
        print("没有需要处理的文件")
        return

    for row in rows:
        process_record(row, db_directory, scan_directory, conn, cursor)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Process files based on photo_move table and move files to the corresponding folders.")
    parser.add_argument('--db', required=True, help="Directory containing the SQLite database file (file_hashes.db)")
    parser.add_argument('--scan-dir', required=True, help="Directory to scan for files")

    args = parser.parse_args()

    db_directory = args.db
    scan_directory = args.scan_dir

    if not os.path.isdir(db_directory):
        print(f"错误: {db_directory} 不是一个有效的目录")
        sys.exit(1)

    if not os.path.isdir(scan_directory):
        print(f"错误: {scan_directory} 不是一个有效的目录")
        sys.exit(1)

    main(db_directory, scan_directory)